import numpy as np
import torch
from .deep.feature_extractor import Extractor
from .sort.nn_matching import NearestNeighborDistanceMetric
from .sort.preprocessing import non_max_suppression
from .sort.detection import Detection
from .sort.tracker import Tracker
import cv2
import configuration as cfg
class DeepSort(object):
    def __init__(self):
        '''
        model_path:deepsort模型权重路径
        max_dist:
        min_confidence:保留的最小置信度
        nms_max_overlap:nms的最大重叠比例
        max_iou_distance:
        '''

        self.extractor = Extractor()   #加载分类器

        # 创建跟踪器，完成对一条轨迹的状态管理、初始化、更新、删除、预测等等。
        self.tracker = Tracker(metric=NearestNeighborDistanceMetric())

    def update(self, bbox_xywh, confidences, ori_img,class_ids):
        '''
        bbox_xywh:bounding box的中心坐标和w,h。
        confidence:置信度
        ori_img:用cv2打开的原始图片，为BGR格式
        '''
        ori_img = cv2.cvtColor(ori_img, cv2.COLOR_BGR2RGB)    #默认的cv2图像是BGR，所以转换成RGB
        self.height, self.width = ori_img.shape[:2]
        #生成检测
        features = self._get_features(bbox_xywh, ori_img)#将ori_img对应用box切片后进行特征提取，是一个向量
        bbox_tlwh = self._xywh_to_tlwh(bbox_xywh)#转换bbox的格式
        #下面这个confidences={Tensor:(2,4)}，第0维表示检测到几个box，并且这个置信度要大于设定的阈值
        '''detections保存了所有的box的信息，每个Detection对象包含了满足阈值的目标置信度，box的特征向量，tlwh的box坐标'''
        detections = [Detection(bbox_tlwh[i], conf, features[i]) for i, conf in enumerate(confidences) if
                      conf > cfg.min_confidence]

        #应用非极大值抑制
        boxes = np.array([d.tlwh for d in detections])#格式为tlwh的所有目标的box坐标信息
        scores = np.array([d.confidence for d in detections])#格式为单个float数的所有box的目标置信度信息
        indices = non_max_suppression(boxes, cfg.nms_max_overlap, scores)  #返回经过NMS剩下boxes的索引
        detections = [detections[i] for i in indices]   #只选取nms过后的索引的detections

        #更新跟踪器
        self.tracker.predict()#track预测
        self.tracker.update(detections,class_ids)#detections与tracks匹配更新track

        # output bbox identities
        outputs = []
        for track in self.tracker.tracks:
            if not track.is_confirmed() or track.time_since_update > 1:
                continue
            box = track.to_tlwh()
            x1, y1, x2, y2 = self._tlwh_to_xyxy(box)
            track_id = track.track_id
            class_id = track.class_id
            outputs.append(np.array([x1, y1, x2, y2, track_id,class_id], dtype=np.int))
        if len(outputs) > 0:
            outputs = np.stack(outputs, axis=0)
        return outputs

    @staticmethod
    def _xywh_to_tlwh(bbox_xywh):
        '''
        转换bbox格式从center_x,center_y,bbox_w,bbox_h
        到top_left_x,top_left_y,bbox_w,bbox_h
        '''
        if isinstance(bbox_xywh, np.ndarray):
            bbox_tlwh = bbox_xywh.copy()
        elif isinstance(bbox_xywh, torch.Tensor):
            bbox_tlwh = bbox_xywh.clone()
        bbox_tlwh[:, 0] = bbox_xywh[:, 0] - bbox_xywh[:, 2] / 2.
        bbox_tlwh[:, 1] = bbox_xywh[:, 1] - bbox_xywh[:, 3] / 2.
        return bbox_tlwh

    def _xywh_to_xyxy(self, bbox_xywh):
        x, y, w, h = bbox_xywh
        x1 = max(int(x - w / 2), 0)
        x2 = min(int(x + w / 2), self.width - 1)
        y1 = max(int(y - h / 2), 0)
        y2 = min(int(y + h / 2), self.height - 1)
        return x1, y1, x2, y2

    def _tlwh_to_xyxy(self, bbox_tlwh):
        """
        TODO:
            Convert bbox from xtl_ytl_w_h to xc_yc_w_h
        Thanks JieChen91@github.com for reporting this bug!
        """
        x, y, w, h = bbox_tlwh
        x1 = max(int(x), 0)
        x2 = min(int(x + w), self.width - 1)
        y1 = max(int(y), 0)
        y2 = min(int(y + h), self.height - 1)
        return x1, y1, x2, y2

    def _get_features(self, bbox_xywh, ori_img):
        '''
        使用box来裁剪ori_img，将裁剪的所有图片boxes作为一个张量。
        将该图片进行特征前向传播，返回结果
        bbox_xywh:预测出来的bbox的center_x,center_y,box_w,box_h，值以原图的高宽为基准
        ori_img:一个cv2打开的ndarry数组，为RGB格式
        '''
        im_crops = []   #用来保存预测出来的box
        for box in bbox_xywh:
            x1, y1, x2, y2 = self._xywh_to_xyxy(box)
            im = ori_img[y1:y2, x1:x2]
            im_crops.append(im)
        if im_crops:
            features = self.extractor(im_crops)
        else:
            features = np.array([])
        return features